// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <stddef.h>
#include <stdint.h>
#include <new>

#include <pollux/common/base/bit_util.h>
#include <pollux/common/base/exceptions.h>

namespace kumo::pollux {
    /// The class provides concurrent updates to a counter with minimum lock
    /// contention. The template argument T specifies the counter type. The counter
    /// is N-way sharded internally. Each update goes to one of the sharded counters
    /// based on the update thread id.
    template<class T>
    class ConcurrentCounter {
    public:
        /// Creates a concurrent counter with specified number of shards.
        ///
        /// NOTE: the constructor sets the actual number of shards to be the next
        /// power of 2.
        explicit ConcurrentCounter(size_t numShards)
            : numShards_(bits::nextPowerOfTwo(numShards)),
              shardMask_(numShards_ - 1),
              counters_(numShards_) {
            POLLUX_CHECK_GE(numShards_, 1);
            for (auto &counter: counters_) {
                counter.value = T();
            }
        }

        ConcurrentCounter(const ConcurrentCounter &) = delete;

        ConcurrentCounter &operator=(const ConcurrentCounter &) = delete;

        /// Invoked to read the sum of values from 'counters_'.
        T read() const {
            T sum = T();
            for (size_t i = 0; i < numShards_; ++i) {
                sum += counters_[i].read();
            }
            return sum;
        }

        /// Invoked to update with 'delta'.
        void update(T delta) {
            counters_[shardIndex()].update(delta);
        }

        /// Invoked to update with 'delta' and user provided 'updateFn'. The function
        /// picks up the shard to apply the customized update.
        using UpdateFn = std::function<bool(T &counter, T delta, std::mutex &lock)>;

        bool update(T delta, const UpdateFn &updateFn) {
            return counters_[shardIndex()].update(delta, updateFn);
        }

        void testingClear() {
            for (auto &counter: counters_) {
                counter.value = T();
            }
        }

        T testingRead(size_t index) const {
            return counters_[index].read();
        }

        bool testingUpdate(size_t index, T delta, const UpdateFn &updateFn) {
            return counters_[index].update(delta, updateFn);
        }

    private:
        struct alignas(melon::hardware_destructive_interference_size) Counter {
            mutable std::mutex lock;
            T value;

            T read() const {
                std::lock_guard<std::mutex> l(lock);
                return value;
            }

            void update(T delta) {
                std::lock_guard<std::mutex> l(lock);
                value += delta;
            }

            bool update(T delta, const UpdateFn &updateFn) {
                return updateFn(value, delta, lock);
            }
        };

        size_t shardIndex() const {
            const size_t hash =
                    std::hash<std::thread::id>{}(std::this_thread::get_id());
            const size_t index = hash & shardMask_;
            POLLUX_DCHECK_LT(index, counters_.size());
            return index;
        }

        const size_t numShards_;
        const size_t shardMask_;

        std::vector<Counter> counters_;
    };
} // namespace kumo::pollux
